import csv


def dataset():
    # 读取 stroke.csv文件
    csvFile = open('./data/stroke.csv', 'r')
    reader = csv.reader(csvFile)

    line = 0
    id, gender, age, hypertension, heart_disease, ever_married, work_type, Residence_type, avg_glucose_level, bmi, smoking_status, stroke = [], [], [], [], [], [], [], [], [], [], [], []

    # 使用原始数据粗略处理，方便查看数据
    for item in reader:
        if line != 0:
            id.append(line)
            gender.append(item[1])
            age.append(item[2])
            hypertension.append(item[3])
            heart_disease.append(item[4])
            ever_married.append(item[5])
            work_type.append(item[6])
            Residence_type.append(item[7])
            avg_glucose_level.append(item[8])
            bmi.append(item[9])
            smoking_status.append(item[10])
            stroke.append(item[11])
        line += 1

    csvFile.close()

    with open("./data/dataset.csv", "wt", encoding="gbk") as f:
        f.write(
            "id,gender,age,hypertension,heart_disease,ever_married,work_type,Residence_type,avg_glucose_level,bmi,smoking_status,stroke\n")
        for i in range(len(gender)):
            f.write("{},{},{},{},{},{},{},{},{},{},{},{}\n".format(str(id[i]), str(gender[i]), str(age[i]),
                                                                   str(hypertension[i]),
                                                                   str(heart_disease[i]), str(ever_married[i]),
                                                                   str(work_type[i]), str(Residence_type[i]),
                                                                   str(avg_glucose_level[i]), str(bmi[i]),
                                                                   str(smoking_status[i]), str(stroke[i])))

    # 标准化数据处理，方便分类处理
    csvFile = open('./data/stroke.csv', 'r')
    reader = csv.reader(csvFile)

    line = 0
    id, gender, age, hypertension, heart_disease, ever_married, work_type, Residence_type, avg_glucose_level, bmi, smoking_status, stroke = [], [], [], [], [], [], [], [], [], [], [], []

    for item in reader:
        if line != 0:
            id.append(line)
            if item[1] == 'Male':
                gender.append("1")
            else:
                gender.append("0")
            age.append(item[2])
            hypertension.append(item[3])
            heart_disease.append(item[4])
            if item[5] == 'Yes':
                ever_married.append("1")
            else:
                ever_married.append("0")
            if item[6] not in work_type:
                work_type.append(item[6])
            if item[7] == 'Urban':
                Residence_type.append("1")
            else:
                Residence_type.append("0")
            avg_glucose_level.append(item[8])
            bmi.append(item[9])
            if item[10] not in smoking_status:
                smoking_status.append(item[10])
            stroke.append(item[11])
        line += 1

    work_type_index, smoking_status_index = [], []
    for i in range(len(work_type)):
        work_type_index.append(i)
    work_type = dict(zip(work_type, work_type_index))
    for i in range(len(smoking_status)):
        smoking_status_index.append(i)
    smoking_status = dict(zip(smoking_status, smoking_status_index))

    print("完成 work_type 分类:\n {} 共 {} 类".format(work_type, len(work_type)))
    print("完成 smoking_status 分类:\n {} 共 {} 类".format(smoking_status, len(smoking_status)))

    csvFile = open('./data/stroke.csv', 'r')
    reader = csv.reader(csvFile)
    line = 0
    new_work_type, new_smoking_status = [], []
    for item in reader:
        if line != 0:
            new_work_type.append(work_type[item[6]])
            new_smoking_status.append(smoking_status[item[10]])
        line = 1
    work_type = new_work_type
    smoking_status = new_smoking_status

    csvFile.close()

    with open("./data/datasets.csv", "wt", encoding="gbk") as f:
        f.write(
            "id,gender,age,hypertension,heart_disease,ever_married,work_type,Residence_type,avg_glucose_level,bmi,smoking_status,stroke\n")
        for i in range(len(gender)):
            f.write("{},{},{},{},{},{},{},{},{},{},{},{}\n".format(str(id[i]), str(gender[i]), str(age[i]),
                                                                   str(hypertension[i]),
                                                                   str(heart_disease[i]), str(ever_married[i]),
                                                                   str(work_type[i]), str(Residence_type[i]),
                                                                   str(avg_glucose_level[i]), str(bmi[i]),
                                                                   str(smoking_status[i]), str(stroke[i])))


if __name__ == '__main__':
    try:
        print("开始进行数据预处理...")
        dataset()
        print("数据预处理完成...")
    except Exception as e:
        print("数据预处理出错:", e)
